import os
import requests
from langchain.llms.base import LLM

### SiliconFlow LLM 实现 ###
API_KEY = os.getenv("CUSTOM_API_KEY", "sk-vdznscbvznobzhmkjmkkvbaymahilxthvhngrxdeczxhjjgt")
BASE_URL = "https://api.siliconflow.cn/v1/chat/completions"

class SiliconFlow(LLM):
    def __init__(self):
        super().__init__()

    @property
    def _llm_type(self) -> str:
        return "siliconflow"

    def siliconflow_completions(self, model: str, prompt: str) -> str:
        """
        调用 SiliconFlow API 生成文本响应
        """
        payload = {
            # "model": "Qwen/Qwen2.5-7B-Instruct",
            "model": "deepseek-ai/DeepSeek-V3",
            "messages": [{"role": "user", "content": prompt}],
            "stream": False  # 关闭流式返回
        }
        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "authorization": f"Bearer {API_KEY}"
        }

        response = requests.post(BASE_URL, json=payload, headers=headers)
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    def _call(self, prompt: str, stop: list = None) -> str:
        """
        返回一个无 Markdown 格式化的纯 SQL 字符串。
        """
        raw_response = self.siliconflow_completions(
            # model="Qwen/Qwen2.5-7B-Instruct",
            model="deepseek-ai/DeepSeek-V3",
            prompt=prompt
        )

        # 处理 Markdown 格式化内容
        if raw_response.startswith("```") and raw_response.endswith("```"):
            raw_response = raw_response.strip("```").strip()  # 去掉 Markdown 格式化部分
            raw_response = raw_response.replace("sql", "")  # 去掉 Markdown 格式化部分

        return raw_response